import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold
import lightgbm as lgbm
import _pickle as pickle
# import time
from lightgbm.sklearn import LGBMClassifier
from sklearn.model_selection import GridSearchCV

# 使用200万行数据进行参数调优
X_train = pd.read_pickle("pkl/X_train")
X_train = X_train.sample(n=2000000, random_state=0, axis=0)
y_train = pd.read_pickle("pkl/y_train")
y_train = y_train.sample(n=2000000, random_state=0, axis=0)

MAX_ROUNDS = 10000
kfold = StratifiedKFold(n_splits=3, shuffle=True, random_state=3)


def get_n_estimators(params, X_train, y_train, early_stopping_rounds=10):
    lgbm_params = params.copy()
    lgbmtrain = lgbm.Dataset(X_train, y_train)
    cv_result = lgbm.cv(lgbm_params, lgbmtrain, num_boost_round=MAX_ROUNDS, nfold=3, metrics='binary_logloss',
                        early_stopping_rounds=early_stopping_rounds, seed=3)
    print('best n_estimators:', len(cv_result['binary_logloss-mean']))
    return len(cv_result['binary_logloss-mean'])


params = {'boosting_type': 'gbdt',
          'objective': 'binary',
          'n_jobs': -1,
          'learning_rate': 0.1,
          'num_leaves': 60,
          'max_depth': 6,
          'max_bin': 127,  # 2^6,原始特征为整数，很少超过100
          'subsample': 0.7,
          'bagging_freq': 1,
          'colsample_bytree': 0.7,
          }

n_estimators_1 = get_n_estimators(params, X_train, y_train)
fo = open("res.txt", "a+")
fo.write("n_estimators_1: " + str(n_estimators_1))
fo.close()
